import os 
import SimpleITK as sitk 
import numpy as np 
import torch.nn as nn 
from ddim_seg.basic_unet_denose import BasicUNetDe
from ddim_seg.basic_unet import BasicUNetEncoder
from guided_diffusion.gaussian_diffusion import get_named_beta_schedule, ModelMeanType, ModelVarType,LossType
from guided_diffusion.respace import SpacedDiffusion, space_timesteps
from guided_diffusion.resample import UniformSampler
import torch 
from monai.inferers import SlidingWindowInferer
from monai import transforms
import matplotlib.pyplot as plt 

def read_data(data_path):
        
    file_identifizer = data_path.split("/")[-1].split("_")[-1]
    image_paths = [
        os.path.join(data_path, f"BraTS20_Training_{file_identifizer}_t1.nii.gz"),
        os.path.join(data_path, f"BraTS20_Training_{file_identifizer}_flair.nii.gz"),
        os.path.join(data_path, f"BraTS20_Training_{file_identifizer}_t2.nii.gz"),
        os.path.join(data_path, f"BraTS20_Training_{file_identifizer}_t1ce.nii.gz")
    ]
    seg_path = os.path.join(data_path, f"BraTS20_Training_{file_identifizer}_seg.nii.gz")

    image_data = [sitk.GetArrayFromImage(sitk.ReadImage(p)) for p in image_paths]
    seg_data = sitk.GetArrayFromImage(sitk.ReadImage(seg_path))

    image_data = np.array(image_data).astype(np.float32)
    seg_data = np.expand_dims(np.array(seg_data).astype(np.int32), axis=0)
    return {
        "image": image_data,
        "label": seg_data
    } 

class DiffUNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.embed_model = BasicUNetEncoder(3, 4, 2, [64, 64, 128, 256, 512, 64])

        self.model = BasicUNetDe(3, 7, 3, [64, 64, 128, 256, 512, 64], 
                                act = ("LeakyReLU", {"negative_slope": 0.1, "inplace": False}))

        betas = get_named_beta_schedule("linear", 1000)

        self.sample_diffusion = SpacedDiffusion(use_timesteps=space_timesteps(1000, [10]),
                                            betas=betas,
                                            model_mean_type=ModelMeanType.START_X,
                                            model_var_type=ModelVarType.FIXED_LARGE,
                                            loss_type=LossType.MSE,
                                            )
        self.sampler = UniformSampler(1000)

    def forward(self, image=None, x=None, pred_type=None, step=None):
        if pred_type == "q_sample":
            noise = torch.randn_like(x).to(x.device)
            t, weight = self.sampler.sample(x.shape[0], x.device)
            return self.diffusion.q_sample(x, t, noise=noise), t, noise

        elif pred_type == "denoise":
            # return self.model(x, t=step, image=image, embed=embed)
            embeddings = self.embed_model(image)
            return self.model(x, t=step, image=image, embeddings=embeddings)

        elif pred_type == "ddim_sample":
            embeddings = self.embed_model(image)

            sample_out = self.sample_diffusion.ddim_sample_loop(self.model, (1, 3, 96, 96, 96), model_kwargs={"image": image, "embeddings": embeddings})
            sample_out = sample_out["pred_xstart"]
            return sample_out

transform = transforms.Compose(
        [   
            transforms.CropForegroundd(keys=["image", "label"], source_key="image"),
            transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            transforms.ToTensord(keys=["image", "label"]),
        ])

def convert_label(output):
    output = output.squeeze(dim=0)
    output = output.cpu().numpy()
    d, w, h = output.shape[1], output.shape[2], output.shape[3]
    show_output = np.zeros((d, w, h))
    wt = output[1] == 1
    tc = output[0] == 1
    et = output[2] == 1
    show_output[wt] = 2
    show_output[tc] = 1
    show_output[et] = 4
    return show_output
    
if __name__ == "__main__":
    device = "cuda:0"
    # 滑窗推理
    window_infer = SlidingWindowInferer(roi_size=[96, 96, 96],
                                        sw_batch_size=1,
                                        overlap=0.5)
    # 新建模型并加载参数
    model = DiffUNet()
    model_path = "./logs_brats/diffusion_seg_all_loss_e300_norm_fix_conv_bugs_embed/model/final_model_0.8696.pt"
    sd = torch.load(model_path)
    model.load_state_dict(sd)
    model.to(device)

    # 读取数据
    input_data = read_data("/home/xingzhaohu/sharefs/datasets/brats2020/MICCAI_BraTS2020_TrainingData/BraTS20_Training_068")

    # 数据预处理
    transform_data = transform(input_data)
    print(f"image shape is {transform_data['image'].shape}")
    print(f"label shape is {transform_data['label'].shape}")
    image = transform_data["image"].unsqueeze(dim=0)
    label = transform_data["label"]
    image = image.to(device)

    # 预测，DDIM采样方式
    print(f"正在预测...")
    with torch.no_grad():
        pred = window_infer(image, model, pred_type="ddim_sample")

    print(f"pred shape is {pred.shape}")
    pred = convert_label(pred)

    # 可视化过程
    image = image.cpu().numpy()
    for i in range(60, 70):
        plt.subplot(2, 4, 1)
        plt.axis("off")
        plt.title("t1")
        plt.imshow(image[0, 0, i], cmap="gray")
        plt.subplot(2, 4, 2)
        plt.axis("off")
        plt.title("flair")
        plt.imshow(image[0, 1, i], cmap="gray")
        plt.subplot(2, 4, 3)
        plt.axis("off")
        plt.title("t2")
        plt.imshow(image[0, 2, i], cmap="gray")
        plt.subplot(2, 4, 4)
        plt.axis("off")
        plt.title("t1ce")
        plt.imshow(image[0, 3, i], cmap="gray")
        plt.subplot(2, 4, 5)
        plt.axis("off")
        plt.title("label")
        plt.imshow(label[0, i], cmap="gray")
        plt.subplot(2, 4, 6)
        plt.axis("off")
        plt.title("pred")
        plt.imshow(pred[i], cmap="gray")
        plt.show()